#!/usr/bin/env python3
# policy/remote_policy_server.py
# 远程策略TCP服务器
import socket
import threading
import time
import struct
import numpy as np

import logging

logger = logging.getLogger(__name__)


def pack_adj_matrix(adj_matrix):
    """将邻接矩阵打包成bit-packed bytes"""
    # 确保是bool类型
    adj_bool = adj_matrix.astype(bool)
    # 计算总bits数
    total_bits = adj_bool.size
    # 计算需要的bytes数
    num_bytes = (total_bits + 7) // 8
    # 展平并打包
    flat_bits = adj_bool.flatten()
    packed = bytearray(num_bytes)
    for i in range(total_bits):
        byte_idx = i // 8
        bit_idx = 7 - (i % 8)  # MSB first
        if flat_bits[i]:
            packed[byte_idx] |= 1 << bit_idx
    return bytes(packed)


def unpack_adj_matrix(packed_bytes, node_num):
    """从bit-packed bytes解包成邻接矩阵"""
    total_bits = node_num * node_num
    num_bytes = len(packed_bytes)
    flat_bits = np.zeros(total_bits, dtype=bool)
    for i in range(total_bits):
        byte_idx = i // 8
        bit_idx = 7 - (i % 8)  # MSB first
        if byte_idx < num_bytes and (packed_bytes[byte_idx] & (1 << bit_idx)):
            flat_bits[i] = True
    return flat_bits.reshape((node_num, node_num)).astype(np.uint8)


class RemotePolicyServer:
    def __init__(self, adj_matrix, node_num, logger=None):
        if logger is None:
            logger = logging.getLogger(__name__)
        self.logger = logger
        self.adj_matrix = adj_matrix
        self.node_num = node_num
        self.policy_socket = None
        self.socket_lock = threading.Lock()
        server_thread = threading.Thread(target=self._start_server, daemon=True)
        server_thread.start()
        self.logger.info("远程策略TCP服务器启动")

    def send_graph_update(self):
        """发送图结构更新包到客户端"""
        with self.socket_lock:
            if self.policy_socket is None:
                self.logger.warning("客户端未连接，无法发送图结构更新")
                return False
            try:
                # 发送包头: 类型 0x01 (图结构更新)
                self.policy_socket.sendall(b"\x01")
                # 发送节点数
                self.policy_socket.sendall(struct.pack(">I", self.node_num))
                # 打包邻接矩阵并发送长度和数据
                adj_packed = pack_adj_matrix(self.adj_matrix)
                self.policy_socket.sendall(struct.pack(">I", len(adj_packed)))
                self.policy_socket.sendall(adj_packed)
                self.logger.info("图结构更新包发送成功")
                return True
            except Exception as e:
                self.logger.error(f"发送图结构更新失败: {e}")
                # 关闭socket并标记为None
                if self.policy_socket:
                    self.policy_socket.close()
                    self.policy_socket = None
                return False

    def _start_server(self):
        """启动TCP服务器"""
        _IP_HOST_ = ("127.0.0.1", 6002)
        self.logger.info(f"TCP服务器监听 {_IP_HOST_}")
        while True:
            time.sleep(1)
            with self.socket_lock:
                if self.policy_socket is not None:
                    try:
                        self.policy_socket.recv(0)
                    except:
                        self.policy_socket.close()
                        self.policy_socket = None
                        self.logger.info("检测到客户端断开，准备重新监听")
                        continue
                    continue

            server_sock = None
            try:
                server_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                server_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
                server_sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
                server_sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 65536)
                server_sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 65536)
                server_sock.settimeout(20.0)
                server_sock.bind(_IP_HOST_)
                server_sock.listen(1)
                with self.socket_lock:
                    self.policy_socket, addr = server_sock.accept()
                self.logger.info(f"远程客户端连接: {addr}")
                # 发送初始图结构更新
                success = self.send_graph_update()
                if not success:
                    raise ConnectionError("初始图结构更新发送失败")
                server_sock.close()
            except Exception as e:
                self.logger.error(f"服务器监听失败: {e}")
                if server_sock:
                    server_sock.close()
                with self.socket_lock:
                    if self.policy_socket:
                        self.policy_socket.close()
                        self.policy_socket = None
                time.sleep(5)

    def calculate_pursuers_move(self, evader_node, p1_node, p2_node):
        """请求远程DRL计算追捕者移动"""
        with self.socket_lock:
            if self.policy_socket is None:
                self.logger.warning("远程客户端未连接，返回当前节点")
                return p1_node, p2_node
            try:
                # 发送包头: 类型 0x02 (策略解算请求)
                self.policy_socket.sendall(b"\x02")
                # 发送数据段: 位置信息
                pack = struct.pack(">III", evader_node, p1_node, p2_node)
                self.policy_socket.sendall(pack)
                self.logger.debug(
                    f"发送策略请求: 逃避者={evader_node}, 追捕者1={p1_node}, 追捕者2={p2_node}"
                )
                # 接收包头: 期望 0x03 (策略解算响应)
                cmd = self.policy_socket.recv(1)
                if not cmd or cmd != b"\x03":
                    raise ValueError(f"意外的响应包头: {cmd}")
                # 接收数据段: 12 字节位置信息
                data = b""
                while len(data) < 12:
                    chunk = self.policy_socket.recv(12 - len(data))
                    if not chunk:
                        raise ConnectionError("远程响应中断")
                    data += chunk
                if len(data) != 12:
                    raise ValueError(f"响应数据不完整: {len(data)} 字节")
                next_evader, p1_next, p2_next = struct.unpack(">III", data)
                self.logger.debug(f"接收响应: p1_next={p1_next}, p2_next={p2_next}")
                return p1_next, p2_next
            except Exception as e:
                self.logger.error(f"远程计算失败: {e}")
                if self.policy_socket:
                    self.policy_socket.close()
                    self.policy_socket = None
                # 重启服务器线程
                threading.Thread(target=self._start_server, daemon=True).start()
                return p1_node, p2_node


if __name__ == "__main__":
    # 生成随机测试数据
    node_count = 10
    # 创建随机邻接矩阵 (无向图)
    adj_matrix_test = np.random.randint(0, 2, (node_count, node_count))
    # 确保矩阵对称（无向图）
    adj_matrix_test = np.maximum(adj_matrix_test, adj_matrix_test.T)
    # 确保对角线为0（无自环）
    np.fill_diagonal(adj_matrix_test, 0)

    print("测试用邻接矩阵:")
    print(adj_matrix_test)

    # 创建远程策略服务器实例
    server = RemotePolicyServer(adj_matrix_test, node_count)

    # 模拟几轮任务步骤
    print("\n开始模拟任务步骤...")
    for step in range(5):
        print(f"\n--- 步骤 {step + 1} ---")

        # 随机生成当前位置
        evader_pos = np.random.randint(0, node_count)
        pursuer1_pos = np.random.randint(0, node_count)
        pursuer2_pos = np.random.randint(0, node_count)

        print(
            f"发送状态: 逃避者={evader_pos}, 追捕者1={pursuer1_pos}, 追捕者2={pursuer2_pos}"
        )

        # 请求计算下一步动作
        p1_next, p2_next = server.calculate_pursuers_move(
            evader_pos, pursuer1_pos, pursuer2_pos
        )

        print(f"接收结果: 追捕者1移动到={p1_next}, 追捕者2移动到={p2_next}")

        # 等待一段时间再进行下一步
        time.sleep(2)

    print("\n测试完成")
